'''
Transform various kinds of input data to Wolfit processes and vice versa.

@author: Philip van Oosten
@license: GNU GPLv2
'''
import wolfit
import networkx as nx
from wolfit import WolfitException
import tempfile
import webbrowser
import time
import os, os.path
import ConfigParser
import textwrap
import sys
import imp
from wolfit.tasks import CallableTask

class WolfitImportException(WolfitException):
    pass
error = WolfitImportException

def display_with_graphviz(process, image_file=None, runnable_process = False, image_type='pdf'):
    if runnable_process:
        process = process.runnable_copy()
    temporary_file = not image_file
    image_file = image_file or tempfile.mktemp('.' + image_type, 'wolfit_test')
    image_type = image_file.split('.')[-1]
    # Display network without the feedback tasks.
    n = process._network
    if runnable_process:
        n = process._network.copy()
        for feedback_task in process.feedback_zones:
            n.remove_node(feedback_task)
    dot_graph = nx.to_pydot(n)
    # Add feedback clusters and feedback edges
    n = process._network
    feedback_zones = process.feedback_zones.values() if runnable_process else ()
    for loop_number, fb_set in enumerate(feedback_zones):
        cluster = nx.pydot.Cluster('feedback' + str(loop_number),
                                   color='blue',
                                   fontcolor='blue',
                                   label='Feedback zone')
        for task in fb_set:
            if task in process.feedback_zones:
                node = nx.pydot.Node(str(task), label='Feedback', shape='triangle',color='maroon',style='filled',fontcolor='white')
            else:
                node = dot_graph.get_node(str(task))
                node.set_parent_graph(cluster)
            cluster.add_node(node)
        for fb_task in process.feedback_zones:
            # Invert connections with successors
            for pred in n.predecessors(fb_task):
                _feedback_dot_edge(dot_graph, fb_task, pred)
            for succ in n.successors_iter(fb_task):
                _feedback_dot_edge(dot_graph, succ, fb_task)
            
            # Invert connections with predecessors
        dot_graph.add_subgraph(cluster)
    dot_graph.write(image_file, format=image_type)
    webbrowser.open_new_tab('file://' + os.path.abspath(image_file))
    if temporary_file:
        time.sleep(5)
        os.unlink(image_file)

class IniFileImporter(object):
    #TODO: avoid code duplication between this class and the wolfit_roundtrip script.
    def _load_ini_config(self, file_name_or_stream):
        config = ConfigParser.SafeConfigParser()
        if isinstance(file_name_or_stream, basestring):
            file_names = config.read(file_name_or_stream)
            if not file_name_or_stream in file_names:
                raise error('Failed to parse process .ini file named {0}'.format(file_name_or_stream))
        else:
            # assume we are dealing with a stream object
            config.readfp(file_name_or_stream)
        return config
        
    def _input_section(self, config, section_name, process, tasks_by_name):
        pass
    def _output_section(self, config, section_name, process, tasks_by_name):
        pass
    def _constants_section(self, config, section_name, process, tasks_by_name):
        pass
    section_handlers = {'input': _input_section,
                        'output': _output_section,
                        'constants': _constants_section,}
    
    def _expect_options(self, config, section, required = (), optional = ()):
        '''
        Checks required and optional options for sections in ini file.
        '''
        options = set(config.options(section))
        required_options_missing = []
        unknown_optional_options = []
        for option in required:
            if not option in options:
                required_options_missing.append(option)
        for option in options:
            if not (option in required or option in optional):
                unknown_optional_options.append(option)
        if required_options_missing:
            raise error('Required options missing in [{0}]: {1}'.format(section, required_options_missing))
        if unknown_optional_options:
            raise error('Unknown options in [{0}]: {1}'.format(section, unknown_optional_options))
    
    def _add_callable_task(self, config, process, section):
        '''
        Adds a python function.
        '''
        self._expect_options(config, section,
                        required = ('type', 'module', 'function'),
                        optional = ('input', 'output', 'before', 'after', 'description'))
        module_name = config.get(section, 'module')
        function_name = config.get(section, 'function')
        if module_name in sys.modules:
            module = sys.modules[module_name]
        else:
            path_module_name = '/'.join(module_name.split('.'))
            module_desc = imp.find_module(path_module_name)
            module = imp.load_module(module_name, *module_desc)
        assert module_name in sys.modules
        function = getattr(module, function_name)
        task = CallableTask(process, function, name=section)
        return process.add_task(task)

    def _add_feedback_task(self, config, process, section):
        self._expect_options(config, section,
                             required=('type', 'condition',),
                             optional=('input', 'output', 'before', 'after', 'description',),)
        condition = config.get(section, 'condition')
        return process.add_feedback(condition = condition)
    
    task_type_handlers = {'python function': _add_callable_task,
                          'feedback': _add_feedback_task,}
    
    def _add_task_from_section(self, config, process, section):
        '''
        Adds a task, the type of which depends on the "type" option.
        '''
        task_type = config.get(section, 'type')
        norm_task_type = self.__norm_name(task_type)
        try:
            handler = self.task_type_handlers[norm_task_type]
        except KeyError:
            raise error(textwrap.dedent('''
            Task in section [{0}] could not be handled.
            Task type "{1}" is not defined.
            '''.format(section, task_type)))
        return handler(self, config, process, section)
    
    def __norm_name(self, section):
        return ' '.join(section.lower().split())
    
    def __multiline_opts(self, option_value):
        for input, output_tuple in [line.split('=') 
                                    for line in [ln.strip()
                                                 for ln in option_value.splitlines()] 
                                    if line]:
            output_task, output = output_tuple.split('.')
            yield [f.strip() for f in input, output_task, output]
    
    def _make_connections(self, tasks_by_name, config):
        for section in config.sections():
            if config.has_option(section, 'input'):
                value = config.get(section, 'input')
                for input, output_task, output in self.__multiline_opts(value):
                    tasks_by_name[section].input[input] = tasks_by_name[output_task].output[output]
            if config.has_option(section, 'output'):
                value = config.get(section, 'output')
                for output, input_task, input in self.__multiline_opts(value):
                    tasks_by_name[input_task].input[input] = tasks_by_name[section].output[output]
            if config.has_option(section, 'before'):
                value = config.get(section, 'before')
                for successor in (t.strip() for t in value.splitlines() if t.strip()):
                    tasks_by_name[section].before(tasks_by_name[successor])
            if config.has_option(section, 'after'):
                value = config.get(section, 'after')
                for predecessor in (t.strip() for t in value.splitlines() if t.strip()):
                    tasks_by_name[section].after(tasks_by_name[predecessor])
        

    def import_ini(self, file_name_or_stream):
        config = self._load_ini_config(file_name_or_stream)
        process = wolfit.Process()
        sections = config.sections()
        section_handlers_to_run = []
        tasks_by_name = {'input': process.inputs_task,
                         'output': process.outputs_task,
                         'constants': process.constants_task}
        # Add tasks
        for section in sections:
            if section.lower() not in self.section_handlers:
                tasks_by_name[section] = self._add_task_from_section(config, process, section)
            else:
                section_handlers_to_run.append(section)
                if section.lower() == 'input':
                    tasks_by_name[section] = process.inputs_task
                elif section.lower() == 'output':
                    tasks_by_name[section] = process.outputs_task
                elif section.lower() == 'constants':
                    tasks_by_name[section] = process.constants_task
        
        # Make connections
        self._make_connections(tasks_by_name, config)
        # Handle sections that are not tasks 
        for section in section_handlers_to_run:
            norm_section = self.__norm_name(section)
            handler = self.section_handlers[norm_section]
            handler(self, config, section, process, tasks_by_name)
        return process

def import_process_from_ini_file(ini_file_or_stream):
    '''
    Creates a process from an ini file.
    '''
    importer = IniFileImporter()
    return importer.import_ini(ini_file_or_stream)

def display_from_ini_file(ini_file_or_stream, *args, **kwargs):
    '''
    Displays the process described in a .ini file as a graphviz layout.
    '''
    process = import_process_from_ini_file(ini_file_or_stream)
    display_with_graphviz(process, *args, **kwargs)
